import math
import sys
import random

from History import *
from Model3 import *
from PathGenerator import *
from GoalGenerator import *

class Controller:
    def __init__(self):
        self.history = History()
        self.goalGenerator = GoalGenerator()
        self.pathGenerator = PathGenerator()
        self.models = []
        self.currentGoalState = None
        self.currentGoalEvent = None
        self.nextState = None
        self.previousState = None
        self.previousModel = None
        self.historyFile = open("his.txt", "w")
        self.tmpFile = open("t.txt", "w")

    def active(self, startEvent, endEvent):
        # TODO: Check if this controller can reach end state, return false if not
        if self.currentGoalState == None:
            self.currentGoalState = endEvent.state
            self.currentGoalEvent = endEvent
            self.planPathToCurrentGoal(startEvent)
            print "Controller: Goal set to: a = " + str(self.currentGoalState.arm1.a) + ", aDot = " + str(self.currentGoalState.arm1.aDot)
        return True
        
    def planPathToCurrentGoal(self, startEvent):
        self.pathGenerator.generatePath(startEvent, self.currentGoalEvent)
        print "Controller: Path set to dAngle = " + str(self.pathGenerator.deltaAngle) + ", dt = " + str(self.pathGenerator.tm)
        
    def searchForGoodStateInAllModels(self):
        for model in self.models:
            model.setGoodStates(self.history.events)
        
    def outsideTrajectory(self, state, time, dt):
        print "Controller: Replanning"
        if self.previousModel:
            self.previousModel.addBadState(self.previousState)
        startEvent = Event(state, time)
        self.planPathToCurrentGoal(startEvent)
        values = (self.fractionOfModelledStatesInHistory(), len(self.history.events), len(self.models))
        print "Controller: Fraction of modeled states in history: {0:5.2f} out of {1} with {2} models".format(values[0], values[1], values[2])
        # TODO: Learn new model
        self.generateModels()
        
    def act(self, state, time, dt):
        print "--------------------------------"
        print "Controller: Acting with input state: {0} @ time {1}".format(state, time)
        self.historyFile.write("{0} {1} {2}\n".format(time, state.arm1.a, state.arm1.aDot))
        #self.actUpdateHistory(state, time, dt)
        # Get next state from path generator
        self.nextState = self.pathGenerator.getNextState(state, time, dt)
        if self.nextState == None:
            print "Controller: No next state in plan. Replanning."
            self.outsideTrajectory(state, time, dt)
            self.nextState = self.pathGenerator.getNextState(state, time, dt)
        if self.nextState == None:
            print "Controller ERROR: Unable to access next state of plan. Planning not possible, but goal not reached."
            self.previousState = state
            return (state.arm1.c1, state.arm1.c2)
        stiffness = self.nextState.arm1.stiffness() 
        return self.actGetControls(state, time, dt, stiffness)
#        return self.actGetControlsPreprogrammed(state, time, dt)
    
    def actUpdateHistory(self, state, time, dt):
        # Set controls of previous state - these were the controls actually used
        if self.previousState:
            self.previousState.arm1.c1 = state.arm1.c1
            self.previousState.arm1.c2 = state.arm1.c2
        # Check if we are close enough to the expected trajectory
        if self.nextState != None:
            print "Controller: Expected state:         ", self.nextState
            print "Controller: Distance to expected state: {0}".format(stateDistanceWithoutControls(state, self.nextState))
            if stateDistanceWithoutControls(state, self.nextState) > 0.5:
                self.outsideTrajectory(state, time, dt)
        # Save state on history stack if sufficiently different from previous state
        if len(self.history.events) > 0:
            prevState = self.history.events[len(self.history.events)-1].state
            prevDt = time - self.history.events[len(self.history.events)-1].time
            #print "Controller: Distance to previous state {0}".format(stateDistance(state, prevState)/prevDt)
            if stateDistance(state, prevState)/prevDt > 10:
                self.history.add(state, time)
        else:
            self.history.add(state, time)
        print "Controller: Size of history = " + str(len(self.history.events))
        # Check if goal was reached
        if self.currentGoalState != None and stateDistance(self.currentGoalState, state) < 0.1:
            print "Goal reached: {0}".format(self.currentGoalState)
            return (state.arm1.c1, state.arm1.c2)

    def actGetControlsPreprogrammed(self, state, time, dt):
        if time < 0.25:
            return (20*time, 0)
        if time < 0.5:
            return (5-20*(time-0.25), 0)
        return (0,0)

    def actGetControls(self, state, time, dt, stiffness):
        # Get the controls necessary to reach next state from model
        print "Controller: State =      {0}".format(state)
        print "Controller: Next state = {0}".format(self.nextState)
        print "Controller: Stiffness  = {0}".format(stiffness)
        model = self.getActiveModel(state, self.nextState, dt, stiffness)
        if model == None:
            print "No active models"
            (c1, c2) = self.guess2(self.previousState, state, self.nextState, stiffness)
            self.previousState = state
            return (c1, c2)
        (c1, c2) = model.getControls(dt, state.arm1.a, state.arm1.aDot, self.nextState.arm1.a, self.nextState.arm1.aDot, stiffness)
        print "Controller: Controls: {0} {1}".format(c1,c2)
        print "Controller: Predicted state: {0}".format(model.nextState(state, dt))
        n = 0
        for m in self.models:
            (c1t, c2t) = m.getControls(dt, state.arm1.a, state.arm1.aDot, self.nextState.arm1.a, self.nextState.arm1.aDot, stiffness)
        #if self.previousState:
        #    a = self.nextState.arm1.aDot - state.arm1.aDot
        #    f = model.k1*(c1-c2) + model.k2*math.cos(state.arm1.a)
            self.tmpFile.write("{0} {1} {2} {3} {4} {5} {6} {7}\n".format(time, n, c1t, c2t, c1, c2, m.k1, m.k2))
            n = n + 1
        self.tmpFile.write("\n")
        # Save state
        self.previousState = state
        self.previousModel = model
        return (c1, c2)
        
    def guess(self, prevState, state, nextState, stiffness):
        print "Controller: Guessing with"
        print "  {0} to".format(prevState)
        print "  {0} to".format(state)
        print "  {0} as input".format(nextState)
        cStep = 0.1
        dc1 = 0
        dc2 = 0
        da = nextState.arm1.a - state.arm1.a
        daDot = nextState.arm1.aDot - state.arm1.aDot
#        print "  da = {0}, aDot = {1}, aDotDot = {2}".format(da, state.arm1.aDot, aDotDot)
        print "  da = {0}, aDot = {1}".format(da, state.arm1.aDot)
        if da > 0:
            dc1 = cStep
        else: # da <= 0:
            dc1 = -cStep
        c1 = state.arm1.c1 + dc1
        c2 = state.arm1.c2
        # Set c2 from new c1 and stiffness
#        if state.arm1.c1 < state.arm1.c2:
#            c2 = (2.0*stiffness+c1) / 3.0
#        else:
#            c2 = 3.0*c1 - 2.0*stiffness
        if c1 < -0.9 or c2 < -0.9:
            c1 += cStep
            c2 += cStep
        return (c1, c2)

    # Use constant speed toward goal state
    def guess2(self, prevState, state, nextState, stiffness):
        print "Controller: Guessing 2 with"
        print "  {0} to".format(prevState)
        print "  {0} to".format(state)
        print "  {0} as input".format(nextState)
        cStep = 0.1
        aDotTarget = 1
        dc1 = 0
        dc2 = 0
        da = self.currentGoalState.arm1.a - state.arm1.a
        print "  da = {0}, aDot = {1}".format(da, state.arm1.aDot)
        if da > 0:
            if state.arm1.aDot > aDotTarget:
                dc1 = -cStep
            else:
                dc1 = cStep
        else: # da <= 0:
            if state.arm1.aDot > -aDotTarget:
                dc1 = -cStep
            else:
                dc1 = cStep
        c1 = state.arm1.c1 + dc1
        c2 = state.arm1.c2
        if c1 < -0.9 or c2 < -0.9:
            c1 += cStep
            c2 += cStep
        return (c1, c2)
        
    def getActiveModel(self, state, nextState, dt, stiffness):
        if len(self.models) == 0:
            return None
        bestModel = self.models[0]
        dGoodBest = bestModel.minDistanceToGoodState(state, nextState, dt, stiffness)
        for model in self.models[1:]:
              dGood = model.minDistanceToGoodState(state, nextState, dt, stiffness)
              if dGood < dGoodBest:
                dGoodBest = dGood
                bestModel = model
        print "Controller: Minimum distance to good state = {0:6.2f} in model {1}".format(dGoodBest, bestModel.id)
        if dGoodBest > 200.0:
            print "Controller: No good model."
            return None
        return bestModel
        
    def generateModels(self, n = 10):
        print "Controller: Generating new models with the {0} newest states in history".format(n)
        start = max(len(self.history.events)-n, 0)
        for n in range(start, len(self.history.events)-1):
            print "Controller: Building models for state {0} {1}".format(n, self.history.events[n].state.modelExists)
            state = self.history.events[n].state
            #print n, state.modelExists
            if n > len(self.history.events)-2:
                break
            if not state.modelExists:
                model = Model3()
                model.setParameters(self.history.events[n:n+3])
                model.setGoodStates(self.history.events[n:])
                #i = n
                #print "model exists for n = {0}: {1}".format(i, self.history.events[i].state.modelExists)
                #i = n+1
                #print "model exists for n = {0}: {1}".format(i, self.history.events[i].state.modelExists)
                #i = n+2
                #print "model exists for n = {0}: {1}".format(i, self.history.events[i].state.modelExists)
                if len(model.goodStates) > 0:
                    print "Controller: Found {0} good states for new model".format(len(model.goodStates))
                    self.models.append(model)
                else:
                    print "Controller: XXXXX No good states for new model"

    def generateModelsLarge(self, n = 10):
        print "Controller: Generating new models with the {0} newest states in history".format(n)
        start = max(len(self.history.events)-n, 0)
        for n in range(start, len(self.history.events)-1):
            print "Controller: Building models for state {0} {1}".format(n, self.history.events[n].state.modelExists)
            state = self.history.events[n].state
            #print n, state.modelExists
            if n > len(self.history.events)-5:
                break
            if not state.modelExists:
                model = Model3()
                model.setParameters(self.history.events[n:n+5])
                model.setGoodStates(self.history.events[n:])
                #i = n
                #print "model exists for n = {0}: {1}".format(i, self.history.events[i].state.modelExists)
                #i = n+1
                #print "model exists for n = {0}: {1}".format(i, self.history.events[i].state.modelExists)
                #i = n+2
                #print "model exists for n = {0}: {1}".format(i, self.history.events[i].state.modelExists)
                if len(model.goodStates) > 0:
                    print "Controller: Found {0} good states for new model".format(len(model.goodStates))
                    self.models.append(model)
                else:
                    print "Controller: XXXXX No good states for new model"
                    
    def printKs(self):
        f = open('ks', 'w')
        for m in self.models:
            for gS in m.goodStates:
                o = State()
                d = stateDistanceWithoutControls(o, gS)
                f.write("{0} {1} {2} {3} {4} {5}\n".format(d, gS.arm1.a, gS.arm1.aDot, m.k1, m.k2, m.k3))
        f.close()

    def fractionOfModelledStatesInHistory(self):
        nModelled = 0
        for event in self.history.events:
            if event.state.modelExists:
                nModelled = nModelled + 1
        return float(nModelled) / float(len(self.history.events))
        
            